"""
# Apriori算法原理（同前，核心改进在数据集生成）
1. 核心思想：基于“频繁项集的子集必频繁”剪枝，迭代生成频繁项集与关联规则
2. 关键指标：
   - 支持度：P(A∩B) = 含A和B的事务数 / 总事务数
   - 置信度：P(B|A) = 含A和B的事务数 / 含A的事务数
3. 数据集改进：新增强关联组合库，提高核心物品共现概率，确保关联规则生成
"""

import pandas as pd
import numpy as np
import random

# ---------------------- 1. 改进版：生成高关联事务数据集 ----------------------
def generate_high_assoc_apriori_dataset(
    num_transactions=1000,  # 总事务数
    base_items=None,        # 基础物品列表
    strong_assoc_pairs=None, # 强关联物品对（核心改进）
    assoc_prob=0.6,         # 事务从强关联组合生成的概率（60%事务优先用关联组合）
    max_items_per_trans=4   # 每个事务最多物品数
):
    """
    生成高关联事务数据集：60%事务从强关联组合生成，40%事务随机生成
    确保“牛奶-面包”“牛奶-鸡蛋”等组合高频出现，易产生关联规则
    """
    # 1. 初始化基础参数（默认物品与强关联组合）
    if base_items is None:
        base_items = ['牛奶', '面包', '鸡蛋', '可乐', '薯片', '饼干', '酸奶', '果汁']
    if strong_assoc_pairs is None:
        # 定义强关联组合（现实中高频共现的物品对/ triples）
        strong_assoc_pairs = [
            {'牛奶', '面包'},
            {'牛奶', '鸡蛋'},
            {'面包', '鸡蛋'},
            {'牛奶', '面包', '鸡蛋'},  # 三元强关联
            {'可乐', '薯片'},          # 零食饮料关联
            {'酸奶', '面包'}           # 早餐关联
        ]
    
    # 2. 调整物品基础权重（核心关联物品权重更高）
    item_weights = {
        '牛奶': 0.9, '面包': 0.8, '鸡蛋': 0.7,  # 核心关联物品（高权重）
        '可乐': 0.4, '薯片': 0.5,              # 次要关联物品
        '饼干': 0.3, '酸奶': 0.6, '果汁': 0.3   # 低权重物品（减少干扰）
    }
    # 权重归一化（用于随机抽样）
    weights = [item_weights[item] for item in base_items]
    weights = np.array(weights) / sum(weights)

    dataset = []
    for _ in range(num_transactions):
        # 3. 60%概率：从强关联组合中生成事务（确保关联组合高频）
        if random.random() < assoc_prob:
            # 随机选择一个强关联组合
            selected_assoc = random.choice(list(strong_assoc_pairs))
            # 可选：在强关联组合基础上，额外随机添加1个低权重物品（增加多样性）
            if len(selected_assoc) < max_items_per_trans and random.random() < 0.3:
                # 从非关联物品中随机选1个（避免破坏强关联）
                non_assoc_items = [item for item in base_items if item not in selected_assoc]
                if non_assoc_items:
                    extra_item = random.choice(non_assoc_items)
                    selected_assoc.add(extra_item)
            # 转换为有序列表（确保项集一致性）
            selected = sorted(list(selected_assoc))
        
        # 4. 40%概率：基于权重随机生成事务（保留部分随机性）
        else:
            # 随机选择事务中的物品数量（1~max_items_per_trans）
            num_items = random.randint(1, max_items_per_trans)
            # 按权重抽样（核心物品更易被选中）
            selected = []
            while len(selected) < num_items:
                item = np.random.choice(base_items, p=weights)
                if item not in selected:
                    selected.append(item)
            selected = sorted(selected)
        
        # 添加到数据集
        dataset.append(','.join(selected))
    
    # 保存为CSV（与代码同级）
    pd.DataFrame({'事务': dataset}).to_csv('apriori_high_assoc_dataset.csv', index_label='事务ID')
    print(f"高关联Apriori数据集生成完成：")
    print(f"- 总事务数：{num_transactions}")
    print(f"- 强关联组合：{[list(pair) for pair in strong_assoc_pairs]}")
    print(f"- 数据集保存路径：apriori_high_assoc_dataset.csv")
    return dataset

# ---------------------- 2. Apriori核心算法（逻辑不变，确保兼容性） ----------------------
def calculate_support(itemset, transactions):
    """计算项集的支持度：含该项目的事务数 / 总事务数"""
    count = 0
    for trans in transactions:
        if itemset.issubset(trans):
            count += 1
    return count / len(transactions)

def generate_c1(transactions):
    """生成1-候选集（所有单个物品，去重排序）"""
    c1 = []
    for trans in transactions:
        for item in trans:
            if not [item] in c1:
                c1.append([item])
    c1.sort()
    return list(map(frozenset, c1))  # frozenset可作为字典键

def generate_ck(frequent_k_1, k):
    """由k-1频繁项集生成k-候选集（自连接：前k-2个元素相同则合并）"""
    ck = []
    len_fk1 = len(frequent_k_1)
    for i in range(len_fk1):
        for j in range(i+1, len_fk1):
            # 取前k-2个元素（确保合并后为k项集）
            fk1_i = list(frequent_k_1[i])[:k-2]
            fk1_j = list(frequent_k_1[j])[:k-2]
            fk1_i.sort()
            fk1_j.sort()
            if fk1_i == fk1_j:
                ck.append(frequent_k_1[i] | frequent_k_1[j])  # 集合合并
    return ck

def generate_frequent_itemset(transactions, min_support=0.2):
    """生成所有频繁项集（支持度≥min_support）"""
    # 步骤1：生成1-频繁项集
    c1 = generate_c1(transactions)
    f1 = []
    support_data = {}  # 存储所有项集的支持度（用于后续规则生成）
    
    for itemset in c1:
        support = calculate_support(itemset, transactions)
        if support >= min_support:
            f1.append(itemset)
        support_data[itemset] = support
    
    frequent_itemsets = [f1]  # 存储所有k-频繁项集
    k = 2
    
    # 步骤2：迭代生成k-频繁项集（k≥2）
    while len(frequent_itemsets[k-2]) > 0:
        ck = generate_ck(frequent_itemsets[k-2], k)
        fk = []
        for itemset in ck:
            support = calculate_support(itemset, transactions)
            if support >= min_support:
                fk.append(itemset)
            support_data[itemset] = support
        frequent_itemsets.append(fk)
        k += 1
    
    # 移除空列表（最后一次迭代可能无新项集）
    frequent_itemsets = [fk for fk in frequent_itemsets if fk]
    return frequent_itemsets, support_data

def generate_association_rules(frequent_itemsets, support_data, min_confidence=0.6):
    """由频繁项集生成关联规则（置信度≥min_confidence）"""
    association_rules = []
    # 仅处理k≥2的频繁项集（1项集无法生成“前件→后件”）
    for k in range(1, len(frequent_itemsets)):
        for itemset in frequent_itemsets[k]:
            # 生成所有非空真子集（作为规则前件）
            subsets = [frozenset([item]) for item in itemset]  # 先从1项前件开始（最常用）
            for subset in subsets:
                consequent = itemset - subset  # 后件 = 项集 - 前件
                # 计算置信度：conf = 支持度(项集) / 支持度(前件)
                support_itemset = support_data[itemset]
                support_subset = support_data[subset]
                confidence = support_itemset / support_subset
                
                if confidence >= min_confidence:
                    association_rules.append({
                        '前件': list(subset),
                        '后件': list(consequent),
                        '支持度': round(support_itemset, 4),
                        '置信度': round(confidence, 4)
                    })
    return association_rules

# ---------------------- 3. 实验运行与结果输出（适配改进数据集） ----------------------
if __name__ == "__main__":
    # 1. 生成高关联数据集（直接运行会自动生成CSV）
    transactions = generate_high_assoc_apriori_dataset(
        num_transactions=1000,  # 1000条事务足够统计关联
        assoc_prob=0.6          # 60%事务来自强关联组合
    )
    
    # 2. 加载数据集（若已生成CSV，可注释上一句，直接加载）
    # df = pd.read_csv('apriori_high_assoc_dataset.csv')
    # transactions = df['事务'].tolist()
    
    # 3. 转换事务格式：字符串→集合（便于计算支持度）
    transactions_set = [set(trans.split(',')) for trans in transactions]
    num_total = len(transactions_set)
    print(f"\n数据集加载完成：共{num_total}条事务")

    # 4. 超参数设置（改进后可适当降低置信度，确保规则生成）
    MIN_SUPPORT = 0.2    # 最小支持度：项集至少出现在20%事务中
    MIN_CONFIDENCE = 0.6 # 最小置信度：规则可信度至少60%（原0.7降低为0.6，更易出结果）

    # 5. 生成频繁项集
    frequent_itemsets, support_data = generate_frequent_itemset(
        transactions_set, min_support=MIN_SUPPORT
    )

    # 6. 生成关联规则
    association_rules = generate_association_rules(
        frequent_itemsets, support_data, min_confidence=MIN_CONFIDENCE
    )

    # ---------------------- 7. 结果打印（清晰展示关联规则） ----------------------
    print(f"\n=== Apriori算法实验结果（高关联数据集） ===")
    print(f"超参数：最小支持度={MIN_SUPPORT}，最小置信度={MIN_CONFIDENCE}")

    # 打印频繁项集
    print(f"\n1. 频繁项集（共{len(frequent_itemsets)}类）：")
    for k, fk in enumerate(frequent_itemsets, start=1):
        print(f"   {k}-频繁项集（{len(fk)}个）：")
        for itemset in fk:
            items = sorted(list(itemset))
            support = support_data[itemset]
            print(f"      {{ {', '.join(items)} }} → 支持度：{support:.4f}")

    # 打印关联规则（核心改进结果）
    print(f"\n2. 关联规则（共{len(association_rules)}条）：")
    if association_rules:
        for i, rule in enumerate(association_rules, start=1):
            antecedent = rule['前件']
            consequent = rule['后件']
            print(f"   规则{i}：{{ {', '.join(antecedent)} }} → {{ {', '.join(consequent)} }}")
            print(f"      支持度：{rule['支持度']}，置信度：{rule['置信度']}")
            # 附加业务解释（帮助理解）
            if {'牛奶'} == set(antecedent) and {'面包'} == set(consequent):
                print(f"      解释：购买牛奶的客户中，{rule['置信度']*100:.1f}%会同时购买面包，该组合出现频率{rule['支持度']*100:.1f}%")
    else:
        # 极端情况：若仍无规则，提示降低阈值
        print(f"   未找到关联规则！可尝试：1. 降低MIN_SUPPORT（如0.15）；2. 降低MIN_CONFIDENCE（如0.5）")
